#!/usr/bin/env python3
"""
nmon MCP Server

面向 nmon 系统监控工具的命令行操作封装。

MCP tools 列表：
1. nmon_file_capture - 执行 nmon 文件捕获模式，支持基本文件输出选项
   参数：filename (str, 可选), capture_count (int, 可选), interval (int, 可选)
   返回：命令执行结果 JSON

2. nmon_advanced_options - 执行 nmon 高级监控选项
   参数：option_type (str), option_value (str, 可选)
   返回：命令执行结果 JSON

3. nmon_batch_mode - 执行 nmon 批处理模式
   参数：capture_count (int, 可选), interval (int, 可选)
   返回：命令执行结果 JSON

4. nmon_system_profiles - 执行 nmon 系统配置文件操作
   参数：profile_type (str)
   返回：命令执行结果 JSON

5. nmon_interactive_controls - 模拟 nmon 交互式控制命令
   参数：control_key (str)
   返回：命令执行结果 JSON
"""

import subprocess
import json
from typing import Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("nmon MCP Server")

@mcp.tool()
def nmon_file_capture(
    filename: Optional[str] = None,
    capture_count: Optional[int] = None,
    interval: Optional[int] = None,
    ctx: Context = None
) -> dict:
    """
    执行 nmon 文件捕获模式，支持基本文件输出选项
    
    参数:
        filename: 输出文件名（可选，使用 -F 选项）
        capture_count: 捕获次数（可选，使用 -c 选项）
        interval: 捕获间隔秒数（可选，使用 -s 选项）
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["nmon", "-f"]
    
    if filename:
        cmd_parts.extend(["-F", filename])
    
    if capture_count:
        cmd_parts.extend(["-c", str(capture_count)])
    
    if interval:
        cmd_parts.extend(["-s", str(interval)])
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Error executing command: {str(e)}"
        }

@mcp.tool()
def nmon_advanced_options(
    option_type: str,
    option_value: Optional[str] = None,
    ctx: Context = None
) -> dict:
    """
    执行 nmon 高级监控选项
    
    参数:
        option_type: 选项类型，如 "a" (GPU/CPU), "d" (磁盘数), "g" (磁盘组), "J" (JFS文件系统), 
                    "l" (每行磁盘数), "m" (输出目录), "M" (MHz), "N" (NFS), "p" (PID), "r" (运行名称)
        option_value: 选项值（可选）
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["nmon", "-f"]
    
    if option_type == "a":
        cmd_parts.append("-a")
    elif option_type == "d" and option_value:
        cmd_parts.extend(["-d", option_value])
    elif option_type == "g" and option_value:
        cmd_parts.extend(["-g", option_value])
    elif option_type == "J":
        cmd_parts.append("-J")
    elif option_type == "l" and option_value:
        cmd_parts.extend(["-l", option_value])
    elif option_type == "m" and option_value:
        cmd_parts.extend(["-m", option_value])
    elif option_type == "M":
        cmd_parts.append("-M")
    elif option_type == "N":
        cmd_parts.append("-N")
    elif option_type == "p":
        cmd_parts.append("-p")
    elif option_type == "r" and option_value:
        cmd_parts.extend(["-r", option_value])
    else:
        return {
            "success": False,
            "command": "",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid option type or missing value: {option_type}"
        }
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Error executing command: {str(e)}"
        }

@mcp.tool()
def nmon_batch_mode(
    capture_count: Optional[int] = None,
    interval: Optional[int] = None,
    ctx: Context = None
) -> dict:
    """
    执行 nmon 批处理模式
    
    参数:
        capture_count: 捕获次数（可选）
        interval: 捕获间隔秒数（可选）
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    cmd_parts = ["nmon", "-b"]
    
    if capture_count:
        cmd_parts.extend(["-c", str(capture_count)])
    
    if interval:
        cmd_parts.extend(["-s", str(interval)])
    
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Error executing command: {str(e)}"
        }

@mcp.tool()
def nmon_system_profiles(
    profile_type: str,
    ctx: Context = None
) -> dict:
    """
    执行 nmon 系统配置文件操作
    
    参数:
        profile_type: 配置文件类型，"z" (系统默认), "x" (扩展), "X" (超级扩展)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    if profile_type not in ["z", "x", "X"]:
        return {
            "success": False,
            "command": "",
            "exit_code": -1,
            "stdout": "",
            "stderr": "Invalid profile type. Must be 'z', 'x', or 'X'"
        }
    
    cmd_parts = ["nmon", f"-{profile_type}"]
    command = " ".join(cmd_parts)
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Error executing command: {str(e)}"
        }

@mcp.tool()
def nmon_interactive_controls(
    control_key: str,
    ctx: Context = None
) -> dict:
    """
    模拟 nmon 交互式控制命令
    
    参数:
        control_key: 控制键，如 "+" (增加间隔), "-" (减少间隔), "0" (大数字), 
                    "1" (进程详情), "3" (内存详情), "c" (CPU), "C" (CPU宽视图), 
                    "d" (磁盘I/O) 等
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    valid_keys = ["+", "-", "0", "1", "3", "c", "C", "d"]
    
    if control_key not in valid_keys:
        return {
            "success": False,
            "command": "",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid control key. Must be one of: {', '.join(valid_keys)}"
        }
    
    # 对于交互式命令，我们使用 echo 来模拟按键输入
    cmd_parts = ["sh", "-c", f"echo '{control_key}' | timeout 5 nmon"]
    command = f"echo '{control_key}' | timeout 5 nmon"
    
    try:
        result = subprocess.run(
            cmd_parts,
            capture_output=True,
            text=True,
            timeout=10
        )
        
        return {
            "success": result.returncode in [0, 124],  # 124 is timeout exit code
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 10 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Error executing command: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()